# -*- coding: utf-8 -*-
"""
PDF 加载器（依赖 BasePdfLoader 通用流程）
仅支持知识库文件的处理方式
"""
import base64
import os
import re
import json
import imghdr
import hashlib
from pathlib import Path
from typing import Any, Dict, List
from urllib.parse import urlparse

from openai import OpenAI
import requests
from loguru import logger

# 兼容服务内导入与单文件运行
try:
    from document_loader.pdf_utils import BasePdfLoader, PdfCommonUtils
except ImportError:
    import sys

    CURRENT_DIR = os.path.dirname(__file__)
    PROJECT_ROOT = os.path.abspath(os.path.join(CURRENT_DIR, os.pardir))
    if PROJECT_ROOT not in sys.path:
        sys.path.append(PROJECT_ROOT)
    from document_loader.pdf_utils import BasePdfLoader, PdfCommonUtils

PDFMetadata = Dict[str, Any]


def cal_md5(file_content):
    file_hash = hashlib.md5()

    file_hash.update(file_content)
    return file_hash.hexdigest()


class KbPdfLoader(BasePdfLoader):
    def load(self, delete_old_file: bool = None) -> PDFMetadata:
        if delete_old_file:
            self.delete_old_file = delete_old_file

        file_pure_name = self.file_path.stem
        md_save_dir = self.base_dir
        md_save_dir.mkdir(parents=True, exist_ok=True)

        parsed_json_path = md_save_dir / f"{file_pure_name}.json"  # 返回结果的json
        parsed_md_path = md_save_dir / f"{file_pure_name}.md"  # 返回结果json中的文本
        cleaned_md_path = md_save_dir / f"{file_pure_name}_cleaned.md"  # 清理过后的md文本
        output_json_path = md_save_dir / f"{file_pure_name}_chunks.json"  # 分句过后的md文本
        result_json_path = md_save_dir / f"{file_pure_name}_results.json"  # 返回结果的json

        # 可以把word转pdf写到这里
        # if self.delete_old_file:
        #     pdf_file_path = md_save_dir / f"{pdf_name}.pdf"
        #     convert_docx_to_pdf()

        if not parsed_md_path.exists():
            logger.debug(f"[Loader] Generating markdown: {parsed_md_path}")
            self._pdf_to_markdown(parsed_json_path, parsed_md_path)

        if self.delete_old_file or not cleaned_md_path.exists():
            logger.debug(f"[Loader] Cleaning markdown: {cleaned_md_path}")
            self._clean_markdown(parsed_md_path, cleaned_md_path)

        if self.delete_old_file or not output_json_path.exists():
            with open(cleaned_md_path, 'r', encoding='utf-8') as f:
                cleaned_text = f.read()
            metadata = self._extract_metadata(cleaned_text, md_filename=file_pure_name)
            result = self._parse_content(cleaned_md_path, output_json_path, metadata)
        else:
            with open(output_json_path, 'r', encoding='utf-8') as f:
                result = json.load(f)

        results = PdfCommonUtils.analyze_result(result)
        if self.delete_old_file or not result_json_path.exists():
            with open(result_json_path, 'w', encoding='utf-8') as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
        return results


class GbPdfLoader(KbPdfLoader):
    def __init__(self, pdf_file_path, delete_old_file: bool = False):
        super().__init__(pdf_file_path, delete_old_file)

        # GB 清洗专属关键段
        self.UNWANTED_SECTIONS = [
            r'目\s*录',
            r'目\s*次',
            r'前\s*言',
            r'规范性引用文件',
            r'参考文献',
            r'引\s*言',
        ]
        self.DATE_PUBLISH_IMPLEMENT_PATTERN = re.compile(
            r'^\s*'
            r'\d{4}[年\-/.]\d{1,2}(?:[月\-/.]\d{1,2}(?:日)?)?'
            r'\s*发\s*布\s*'
            r'\d{4}[年\-/.]\d{1,2}(?:[月\-/.]\d{1,2}(?:日)?)?'
            r'\s*实\s*施\s*'
            r'(?:\r?\n)?',
            flags=re.MULTILINE
        )

        self.API_URL = "https://api.textin.com/ai/service/v1/pdf_to_markdown"
        self.headers = {
            "x-ti-app-id": "8e5ac65e5c09273415e4d6275493e3c3",
            "x-ti-secret-code": "f240bfbfcebbb42f9be0dafce5c0f44d",
            "Content-Type": "application/octet-stream",
        }

    # ----- Base hooks -----
    def _pdf_to_markdown(self, json_path: Path, md_path: Path) -> None:
        """调用 TextIn API 将 PDF 转为 Markdown，并存储原始结果 json。"""
        try:
            params = {
                "page_start": 0,
                "page_count": 1000,
                "parse_mode": "scan",
                "table_flavor": "html",
                "apply_document_tree": 1,
                "page_details": 1,
                "markdown_details": 1,
                "get_image": "objects",
                "image_output_type": "default",
                "raw_ocr": 1,
            }
            file_data = self.file_path.read_bytes()
            resp = requests.post(self.API_URL, headers=self.headers, params=params, data=file_data)
            resp.raise_for_status()
            j = resp.json()
            if j.get("code") != 200:
                raise RuntimeError(f"API 错误: {j.get('message')}")
            md_result = j["result"]
            json_path.write_text(json.dumps(md_result, ensure_ascii=False, indent=2), encoding="utf-8")
            md_path.write_text(md_result.get("markdown", ""), encoding="utf-8")
            logger.debug(f"✅ 已保存：{json_path} 和 {md_path}")

        except Exception as e:
            logger.error(f"PDF转换失败: {e}")
            raise

    def _clean_markdown(self, in_path: Path, out_path: Path) -> None:
        """GB 专属清洗规则：移除部分章节块"""

        def make_section_pattern(section_name: str):
            header_regex = rf'^(?P<hashes>#+)\s*(?:\d+(?:\.\d+)?\s*)?{section_name}.*\r?\n'
            body_regex = r'(?:.*\r?\n?)*?'
            end_regex = r'(?=^#{1,}\s|\Z)'
            return re.compile(header_regex + body_regex + end_regex, flags=re.MULTILINE)

        content = in_path.read_text(encoding='utf-8')
        content = self.DATE_PUBLISH_IMPLEMENT_PATTERN.sub('', content)
        for section in self.UNWANTED_SECTIONS:
            content = make_section_pattern(section).sub('', content)
        out_path.write_text(content, encoding='utf-8')
        logger.debug(f"[GbPdfLoader] Cleaned markdown saved: {out_path}")
    
    def _extract_metadata(self, cleaned_md_text: str, **kwargs) -> Dict[str, Any]:
        def extract_gb_number(content: str) -> str:
            for line in content.split('\n')[:20]:
                m = re.search(r'GB[/T\s]*\s*\d+(?:\.\d+)*-\d{4}', line.strip())
                if m:
                    return m.group().strip()
            return ""

        def extract_scope(content: str) -> str:
            lines = content.split('\n')
            scope_start = -1
            for i, line in enumerate(lines):
                line = line.strip()
                if re.match(r'##\s*1\s*范围', line):
                    scope_start = i + 1
                    break
            if scope_start == -1:
                return ""
            scope_end = -1
            for i in range(scope_start, len(lines)):
                line = lines[i].strip()
                if re.match(r'##\s*\d+', line) and not re.match(r'##\s*1\s*范围', line):
                    scope_end = i
                    break
            if scope_end == -1:
                scope_end = len(lines)
            return '\n'.join(lines[scope_start:scope_end]).strip() or ""

        def extract_file_name(content: str):
            """从markdown内容中提取文件名"""
            # 去掉空行
            content = '\n'.join([line.strip() for line in content.split('\n') if line.strip()])
            lines = content.split('\n')

            # 收集前20行中以"#"开头的行
            title_lines = []
            for i, line in enumerate(lines[:20]):
                stripped_line = line.strip()
                # 匹配以一个或多个"#"开头的行（后面可以有空格也可以没有）
                if re.match(r'^#+', stripped_line):
                    # 提取标题内容，去掉所有的"#"和前后空格
                    title_content = re.sub(r'^#+\s*', '', stripped_line).strip()
                    if title_content:  # 确保提取出的内容不为空
                        title_lines.append(title_content)

            # 查找相同的标题内容
            for i in range(len(title_lines)):
                for j in range(i + 1, len(title_lines)):
                    if title_lines[i] == title_lines[j] and title_lines[i]:
                        return title_lines[i]

            return ''

        gb_number = extract_gb_number(cleaned_md_text)
        return {
            "gb_number": gb_number,
            "scope": extract_scope(cleaned_md_text),
            "filename": extract_file_name(cleaned_md_text) + ' ' + gb_number
        }


class GbPdfLoaderV2(GbPdfLoader):
    def _extract_metadata(self, cleaned_md_text: str, md_filename: str, **kwargs) -> Dict[str, Any]:
        """
        使用qwen3模型一次性提取文件名、国标编号、使用范围、要素分类
        参数：
            content: 文档的markdown内容
            md_filename: markdown文件本身的文件名（可作为参考）
        返回：
            包含四个字段的字典：{"filename": "", "gb_number": "", "scope": "", "category": ""}
        """
        QWEN3_API_KEY = "sk-6f89326783dd42b3b9c63284ef454cf2"
        QWEN3_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
        QWEN3_MODEL_ID = "qwen3-32b"

        # 只取前2000个字符
        preview_content = cleaned_md_text[:2000]
        
        prompt = f"""
            请从以下markdown文档的开头部分中提取四个信息：文件名、国标编号、使用范围、要素分类。

            【文档的文件名（可作为参考）】
            {md_filename}

            【文档开头内容（前2000字符）】
            {preview_content}

            【提取要求】
            1. **文件名**：
            - 通常文档名称会在开头部分出现（作为标题），一般是以"#"开头的标题内容；
            - 如果是国家标准数据，返回的文件名需要是："文件名 GB编号"，返回结果例如："呼吸防护 自吸过滤式防颗粒物呼吸器 GB 2626-2019"等；
            - 如果不是国家标准数据，返回的文件名需要是："文件名"，参考【文档的文件名】，返回结果例如："《中华任命共和国港口法》中华人民共和国主席令第二十三号，2018年12月29日起施行"等；
            - 严格按照上述要求返回，不要包含任何其他内容；

            2. **国标编号**：
            - 格式通常为：GB xxxx-xxx；例如：GB 2626-2019；
            - 国标编号一般在文档开头部分出现，需要根据文档内容判断；
            - 只有国标数据才有国标编号，其他数据没有国标编号；
            - 国标数据是用在返回的【文件名】中的；

            3. **使用范围**：
            - 如果是国家标准数据，相关内容通常在“范围”、“适用范围”、“适用范围”章节中；
            - 如果不是国家标准数据，相关内容通常会表述为：“xxx，适用于本规定、条例、办法等”或类似的描述；
            - 如果无相关内容，返回空字符串：""，不要包含任何其他内容；
            
            4. **要素分类**：
                - 根据文档内容，给出要素分类，必须有值，且只允许有一个值；
                - 需要根据文档内容判断，返回结果数据类型为字符串；
                - 可选的值有：
                    - 设计
                    - 供电、照明
                    - 防雷、防静电
                    - 消防
                    - 安全设施、附属设施
                    - 个体防护设施
                    - 危险货物
                    - 危险化学品
                    - 风险辨识与评估
                    - 安全规程
                    - 通航
                    - 港口设施保安
                    - 应急管理
                    - 特种设备

            【输出要求】
            必须严格按照以下JSON格式输出，不要包含任何其他内容：
            ```json
            {{
            "filename": "提取到的文件名",
            "gb_number": "提取到的国标编号",
            "scope": "提取到的使用范围内容",
            "category": "xx"
            }}
            ```
        """
        
        try:
            client = OpenAI(api_key=QWEN3_API_KEY, base_url=QWEN3_BASE_URL)
            
            response = client.chat.completions.create(
                model=QWEN3_MODEL_ID,
                messages=[
                    {"role": "system", "content": "你是一个专业的文档信息提取专家，擅长从给定的文档中准确提取关键信息。"},
                    {"role": "user", "content": prompt}
                ],
                stream=False,
                extra_body={"enable_thinking": False},
                response_format={"type": "json_object"}
            )
            
            result_text = response.choices[0].message.content.strip()
            
            # 解析JSON
            result_json = json.loads(result_text)
            
            # print(f"模型提取成功:")
            # print(f"  - 文件名: {result_json.get('filename', '未提取到')}")
            # print(f"  - 国标编号: {result_json.get('gb_number', '未提取到')}")
            # print(f"  - 使用范围: {result_json.get('scope', '未提取到')[:100]}...")
            # print(f"  - 要素分类: {result_json.get('category', '未提取到')}")

            return result_json
            
        except Exception as e:
            print(f"调用qwen3 API失败: {e}")
            return {"filename": "", "gb_number": "", "scope": "", "category": ""}


class CommonPdfLoader(KbPdfLoader):
    """通用PDF加载器：使用本地PPStructureV3将PDF转为Markdown，其余流程由BasePdfLoader统一处理"""

    def __init__(self, pdf_file_path: str, delete_old_file: bool = False):
        super().__init__(pdf_file_path, delete_old_file)
        self.UNWANTED_SECTIONS = [
            r'目\s*录',
            r'目\s*次',
            r'前\s*言',
            r'规范性引用文件',
            r'参考文献',
            r'引\s*言',
            r'目的',
            r'适用范围',
            r'使用范围',
            r'作业前准备',
        ]
        self.API_URL = "http://60.205.246.174:8999/service/v1/pdf_to_markdown"
        self.headers = {
            "x-ti-app-id": "8e5ac65e5c09273415e4d6275493e3c3",
            "x-ti-secret-code": "f240bfbfcebbb42f9be0dafce5c0f44d",
        }
        # 预编译
        self._compiled_section_patterns = [self._make_section_pattern(s) for s in self.UNWANTED_SECTIONS]

    def _make_section_pattern(self, section_name: str) -> re.Pattern:
        # 匹配标题（允许前缀编号）到下一个标题/文件末尾
        header = rf'^(?P<h>#+)\s*(?:\d+(?:\.\d+)*[.\s]*)?{re.escape(section_name)}.*\r?\n'
        body = r'(?:.*\r?\n?)*?'
        end = r'(?=^#{1,}\s|\Z)'
        return re.compile(header + body + end, flags=re.MULTILINE)

    def _pdf_to_markdown(self, json_path: Path, md_path: Path) -> None:
        try:
            file_data = self.file_path.read_bytes()
            resp = requests.post(
                self.API_URL,
                headers=self.headers,
                files=[
                    ('file', (self.file_path.name, file_data, 'application/pdf'))
                ]
            )
            resp.raise_for_status()

            j = resp.json()
            if j.get("code") != 200:
                raise RuntimeError(f"API 错误: {j.get('message')}")

            json_path.write_text(json.dumps(j["data"]["result"], ensure_ascii=False, indent=2), encoding="utf-8")
            md_path.write_text(j["data"]["result"].get("markdown", ""), encoding="utf-8")

        except Exception as e:
            logger.error(f"PDF转换失败: {e}")
            raise

    def _clean_markdown(self, in_path: Path, out_path: Path) -> None:
        def clean_content(text: str) -> str:
            for p in self._compiled_section_patterns:
                text = p.sub('', text)
            return text

        content = in_path.read_text(encoding='utf-8')
        cleaned = clean_content(content)
        out_path.write_text(cleaned, encoding='utf-8')
        logger.debug(f"[CommonPdfLoader] Cleaned markdown saved: {out_path}")


if __name__ == "__main__":
    # gb文件
    gb_pdf_path = '/Users/wzq/Desktop/portAI/yoka-main/document_loader/test_data/222/222.md'
    # loader = GbPdfLoader(gb_pdf_path, delete_old_file=False)
    loader = GbPdfLoaderV2(gb_pdf_path, delete_old_file=False)
    result = loader.load()
    # 仅输出关键信息，详细结果已存储为同名 *_chunks.json
    print(json.dumps(result, ensure_ascii=False, indent=2))

    # 非gb文件
    # other_pdf_path = "/document_loader/test_data/other-file/TSPGC.GS.WJ20_32安全生产应急管理办法.pdf"
    # pdf_loader = CommonPdfLoader(other_pdf_path, delete_old_file=False)
    # result = loader.load()
    # # 仅输出关键信息，详细结果已存储为同名 *_chunks.json
    # print(json.dumps(result, ensure_ascii=False, indent=2))

